import torchvision
import torchvision.transforms as transforms
import numpy as np

def load_mnist():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    return trainset, testset

def load_cifar10():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    return trainset, testset

def dataset_to_numpy(dataset, flatten=True):
    X = []
    y = []
    for img, label in dataset:
        arr = img.numpy()
        if flatten:
            arr = arr.flatten()
        else:
            arr = arr.transpose(1, 2, 0)  # C,H,W -> H,W,C
        X.append(arr)
        y.append(label)
    X = np.array(X)
    y = np.array(y)
    return X, y
